import cst.interface
import cst.results
import matplotlib.pyplot as plt
import numpy as np
import os

"""input"""
filename = r"Your\own\path\5x5_1.cst"
model_name = 'cnn_sim2024-11-25_09-09'
freq_num = 151
test_size = 5
"""input end"""

interface = cst.interface.DesignEnvironment()
project = interface.open_project(filename)
modeler = project.modeler


def change_material(component, name, material):
    return 'Solid.ChangeMaterial "%s:%s", "%s"' % (component, name, material)


def init_grid():
    command_list = [change_material('component2', str(block_i), 'bac') for block_i in range(25)]
    for block_i in range(25):
        modeler.add_to_history('change material: component2:%d to: bac' % block_i, command_list[block_i])

    return


def build_grid(grid_config):
    for block_i in range(25):
        if grid_config[block_i] == 1:
            modeler.add_to_history('change material: component2:%d to: bac' % block_i,
                                   change_material('component2', str(block_i), 'Al_drude'))
    return


input_file_name = f'test/test_grid_x{test_size}_{model_name}.npy'
model_output_file_name = f'test/test_out_x{test_size}_{model_name}.npy'
input_array = np.load(input_file_name)
model_output_array = np.load(model_output_file_name)
sample_num = input_array.shape[0]
res_array = np.zeros((sample_num, freq_num * 4 * 2))

for i in range(sample_num):
    init_grid()
    grid = input_array[i]
    build_grid(grid)
    modeler.run_solver()
    results_3d = cst.results.ProjectFile(filename, allow_interactive=True).get_3d()
    tyx = results_3d.get_result_item(r"Tables\1D Results\SZmax(1),Zmin(2)_1D_xResample")
    txx = results_3d.get_result_item(r"Tables\1D Results\SZmax(2),Zmin(2)_1D_xResample")
    ryx = results_3d.get_result_item(r"Tables\1D Results\SZmin(1),Zmin(2)_1D_xResample")
    rxx = results_3d.get_result_item(r"Tables\1D Results\SZmin(2),Zmin(2)_1D_xResample")
    tyx_real = np.real(np.array(tyx.get_ydata()))
    tyx_imag = np.imag(np.array(tyx.get_ydata()))
    txx_real = np.real(np.array(txx.get_ydata()))
    txx_imag = np.imag(np.array(txx.get_ydata()))
    ryx_real = np.real(np.array(ryx.get_ydata()))
    ryx_imag = np.imag(np.array(ryx.get_ydata()))
    rxx_real = np.real(np.array(rxx.get_ydata()))
    rxx_imag = np.imag(np.array(rxx.get_ydata()))
    res_array[i, :] = np.concatenate((tyx_real, tyx_imag, txx_real, txx_imag, ryx_real, ryx_imag, rxx_real, rxx_imag))

np.save(f'test/CST_out_x{test_size}_{model_name}.npy', res_array)

freq_list = np.linspace(150, 300, freq_num)
s_para_name_list = ['tyx', 'txx', 'ryx', 'rxx']

try:
    os.makedirs(f'test/fig/{model_name}')
except FileExistsError:
    pass

for sample_idx in range(sample_num):
    sample_fig_path = f'test/fig/{model_name}/sample{sample_idx}'

    try:
        os.makedirs(sample_fig_path)
    except FileExistsError:
        pass

    for s_para_idx in range(4):
        plt.figure(figsize=(8, 4), dpi=200)

        plt.plot(freq_list,
                 model_output_array[sample_idx, s_para_idx * freq_num * 2: (s_para_idx * 2 + 1) * freq_num],
                 '-', label='NN_real')
        plt.plot(freq_list,
                 model_output_array[sample_idx, s_para_idx * freq_num * 2 + freq_num: (s_para_idx * 2 + 2) * freq_num],
                 '-', label='NN_imag')
        plt.plot(freq_list,
                 res_array[sample_idx, s_para_idx * freq_num * 2: (s_para_idx * 2 + 1) * freq_num],
                 '--', label='CST_real')
        plt.plot(freq_list,
                 res_array[sample_idx, s_para_idx * freq_num * 2 + freq_num: (s_para_idx * 2 + 2) * freq_num],
                 '--', label='CST_imag')
        plt.legend()
        plt.xlabel('Freq(THz)')
        plt.ylabel(s_para_name_list[s_para_idx])
        plt.tight_layout()
        plt.savefig(f'{sample_fig_path}/{s_para_name_list[s_para_idx]}.png')
